from sqlalchemy import select

from owl_admin.ext import db
from owl_ai.domain.ai_model_entity import LLMModel
from owl_ai.domain.ai_model_po import LLMModelPO
from owl_common.sqlalchemy.model import ColumnEntityList


class ChatModelMapper:

    default_fields = ['id', 'model_name', 'api_type', 'api_host', 'api_key', 'vision_support', 'model_describe', 'audio_support', 'create_by', 'create_time', 'update_by', 'update_time']

    default_columns = ColumnEntityList(LLMModelPO, default_fields, False)

    @classmethod
    def find_model_list(cls):
        criteria = []

        stmt = select(*cls.default_columns).where(*criteria).order_by(LLMModelPO.update_time.desc()).limit(10)
        rows = db.session.execute(stmt).all()
        return [cls.default_columns.cast(row, LLMModel) for row in rows]

    @classmethod
    def find_model_by_id(cls, model_id: int):
        criteria = [LLMModelPO.id == model_id]
        stmt = select(*cls.default_columns).where(*criteria)
        row = db.session.execute(stmt).one_or_none()
        return cls.default_columns.cast(row, LLMModel) if row else None